package com.yc.config;

import com.clickhouse.data.ClickHouseDataType;
import com.yc.service.ToolServices;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.community.store.embedding.clickhouse.ClickHouseEmbeddingStore;
import dev.langchain4j.community.store.embedding.clickhouse.ClickHouseSettings;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.*;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;

@Configuration
public class AiConfig {

    // 高级API保持不变
    public interface Assistant {
        @SystemMessage("你是一位专业的老师,要帮忙查询学生的成绩和考试情况")
        String chat(@MemoryId int memoryId, @UserMessage String question);

        @SystemMessage("你是一位电商网站的客服")
        TokenStream chatStream(@MemoryId int memoryId, @UserMessage String question);


    }

    // 元数据类型映射
    private Map<String, ClickHouseDataType> createMetadataTypeMap() {
        Map<String, ClickHouseDataType> metadataTypeMap = new HashMap<>();
        // 添加通用的元数据字段类型
        return metadataTypeMap;
    }

    // 创建不同表的元数据类型映射
    private Map<String, ClickHouseDataType> createPdfMetadataTypeMap() {
        Map<String, ClickHouseDataType> map = createMetadataTypeMap();
        map.put("page_number", ClickHouseDataType.Int32);
        map.put("author", ClickHouseDataType.String);
        map.put("title", ClickHouseDataType.String);
        return map;
    }
    private Map<String, ClickHouseDataType> createWordMetadataTypeMap() {
        Map<String, ClickHouseDataType> map = createMetadataTypeMap();
        map.put("paragraph_number", ClickHouseDataType.Int32);
        map.put("style", ClickHouseDataType.String);
        map.put("is_heading", ClickHouseDataType.UInt8);
        return map;
    }

    private Map<String, ClickHouseDataType> createExcelMetadataTypeMap() {
        Map<String, ClickHouseDataType> map = createMetadataTypeMap();
        map.put("sheet_name", ClickHouseDataType.String);
        map.put("row_number", ClickHouseDataType.Int32);
        map.put("col_number", ClickHouseDataType.Int32);
        return map;
    }

    @Bean
    public EmbeddingModel embeddingModel() {
        return QwenEmbeddingModel.builder()
                .apiKey(System.getenv("ALI_API_KEY"))
                .build();
    }

    // 基础ClickHouse配置
    private ClickHouseSettings createSettings(String tableName, Map<String, ClickHouseDataType> metadataTypeMap) {
        return ClickHouseSettings.builder()
                .url("http://localhost:8123")
                .table(tableName)
                .username("click")
                .password("click123")
                .dimension(embeddingModel().dimension())
                .metadataTypeMap(metadataTypeMap)
                .build();
    }

    // 默认向量存储
    @Bean(name = "defaultEmbeddingStore")
    public ClickHouseEmbeddingStore defaultEmbeddingStore() {
        return ClickHouseEmbeddingStore.builder()
                .settings(createSettings("default_embeddings", createMetadataTypeMap()))
                .build();
    }

    // PDF向量存储
    @Bean(name = "pdfEmbeddingStore")
    public ClickHouseEmbeddingStore pdfEmbeddingStore() {
        return ClickHouseEmbeddingStore.builder()
                .settings(createSettings("pdf_embeddings", createPdfMetadataTypeMap()))
                .build();
    }

    // Word向量存储
    @Bean(name = "wordEmbeddingStore")
    public ClickHouseEmbeddingStore wordEmbeddingStore() {
        return ClickHouseEmbeddingStore.builder()
                .settings(createSettings("word_embeddings", createWordMetadataTypeMap()))
                .build();
    }

    // Excel向量存储
    @Bean(name = "excelEmbeddingStore")
    public ClickHouseEmbeddingStore excelEmbeddingStore() {
        return ClickHouseEmbeddingStore.builder()
                .settings(createSettings("excel_embeddings", createExcelMetadataTypeMap()))
                .build();
    }

    // 生成代理对象
    @Bean
    public Assistant assistant(
            ChatModel chatModel,
            StreamingChatModel streamingChatModel,
            ChatMemoryStore chatMemoryStore,
            ToolServices tools,
            @Qualifier("excelEmbeddingStore") ClickHouseEmbeddingStore excelEmbeddingStore,
            @Qualifier("excelEmbeddingStore") ClickHouseEmbeddingStore embeddingStore,
            QwenEmbeddingModel qwenEmbeddingModel) {

        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(1000)
                .chatMemoryStore(chatMemoryStore)
                .build();

        EmbeddingStoreContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(qwenEmbeddingModel)
                .embeddingStore(embeddingStore)
                .build();

        // 为工具服务配置正确的向量库
        tools.setScoreEmbeddingStore(excelEmbeddingStore); // 确保工具服务使用正确的向量库

        return AiServices.builder(Assistant.class)
                .chatModel(chatModel)
                .streamingChatModel(streamingChatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .contentRetriever(retriever)
                .tools(tools)
                .build();
    }
}